package websocket

import (
	"github.com/star-find-cloud/star-chat/ErrorCode"
	"sync"
)

type ConnectManager struct {
	connections map[int64]*Connect // userID -> Connect
	mutex       sync.RWMutex
}

var (
	manager *ConnectManager
	once    sync.Once
)

func GetManager() *ConnectManager {
	once.Do(func() {
		manager = &ConnectManager{
			connections: make(map[int64]*Connect),
		}
	})
	return manager
}

// Register 注册用户连接
func (cm *ConnectManager) Register(userID int64, conn *Connect) {
	cm.mutex.Lock()
	defer cm.mutex.Unlock()
	cm.connections[userID] = conn
}

// Unregister 移除用户连接
func (cm *ConnectManager) Unregister(userID int64) {
	cm.mutex.Lock()
	defer cm.mutex.Unlock()
	if conn, exists := cm.connections[userID]; exists {
		close(conn.Send)
		delete(cm.connections, userID)
	}
}

// GetConnection 获取用户连接
func (cm *ConnectManager) GetConnection(userID int64) (*Connect, bool) {
	cm.mutex.RLock()
	defer cm.mutex.RUnlock()
	conn, exists := cm.connections[userID]
	return conn, exists
}

// SendToUser 向指定用户发送消息
func (cm *ConnectManager) SendToUser(userID int64, message *Message) error {
	cm.mutex.RLock()
	defer cm.mutex.RUnlock()

	if conn, exists := cm.connections[userID]; exists {
		select {
		case conn.Send <- message:
			return nil
		default:
			return ErrorCode.ERROR_Code_ChannelFull
		}
	}
	return ErrorCode.ERROR_Code_WS_UserNotConnected
}

// Broadcast 广播消息给所有用户
func (cm *ConnectManager) Broadcast(message *Message) {
	cm.mutex.RLock()
	defer cm.mutex.RUnlock()

	for userID, conn := range cm.connections {
		select {
		case conn.Send <- message:
			// 发送成功
		default:
			// 处理通道满的情况，可以选择移除该连接
			go cm.Unregister(userID)
		}
	}
}
